# 要添加一个新单元，输入 '# %%'
# 要添加一个新的标记单元，输入 '# %% [markdown]'
# 我发现写到最后，还是学弟写的那种样子。好麻烦啊，感觉应该一天就能写完，但自己在硬拖。
import copy
import torch
from torch.utils import data
from trainers.fedbase.FedClientBase import FedClientBase

import logging
# logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logging.basicConfig(format ='%(message)s')
logger = logging.getLogger()
logger.setLevel(logging.INFO)


class FedPerClient(FedClientBase):
    """训练相关的参数本来应该是服务器发送过来的，包括模型、参数啥的。训练数据应该是本地的。\n
    这里为了方便，只模拟联邦学习的效果和逻辑，不模拟联邦学习的真是过程。
    1. 数据由服务器分成独立同分布，或者非独立同分布。然后分配给各个客户端。而非客户端在本地读取。
    2. 训练的模型、参数在客户端配置。
    3. 客户端完成训练过程、服务端完成聚合过程。客户端和服务端只发送和传输模型和优化器的权重参数。完成聚合和分配的通信逻辑。

    接收参数，完成训练。主要是一次训练的参数。
    每一个算法都应该提供梯度的返回值。即在客户端计算出梯度。其实在服务器计算也无可厚非。

    Args:
        FedClientBase ([type]): [description]

    Returns:
        [type]: [description]
    """    

    def __init__(self,datapair,model):
        super(FedPerClient, self).__init__(datapair,model)    
        logger.info("FedPerClient init----")



    def set_parameters(self,client_params):
        """从服务器加载训练的参数。并根据batch_size等参数划分数据集。
        应该对以下内容进行区分。在联邦个性化层中。包括base_layer和personalization_layer.
        元学习中，并没有将模型进行分层。所以这里不能命名为meta layer

        Args:
            client_params (dict): 客户端进行训练的相关参数
        """     
        # 训练基本参数(包含记载数据)---------------------------------------------------------------------------------------
        super().set_parameters(client_params)   
        # FedAmpClient策略参数-------------------------------------------------------------------------
        # fedper数据集划分参数
        if("support_ratio" in client_params):
            self.support_ratio=client_params["support_ratio"]
        else:
            self.support_ratio = 0.7

        # fedper个性化层的参数
        if("base_layer" in client_params):
            self.base_layer=client_params["base_layer"]
        else:
            self.base_layer = 4


        logger.info("FedPerClient set_parameters:support_ratio={}\t base_layer={}".format(self.support_ratio,self.base_layer))

    
    def process_data(self):
        """标准的数据处理过程。调用了self.process_data_support_and_query

        两种数据划分方式都实现，因为在测试的时候需要划分两种数据集。
        将其训练方法，不再与元学习对标（元学习训练训练端和测试客户端分离，在测试客户端上进行一次本地适应再学习）。
        而是与amp对标（amp训练的训练端与测试端为一体的。在worker上训练完后，用worker上的测试数据本地测试）。
        """        
        # x,y = self.datapair
        # assert len(x) == len(y)
        # # 将数据转换成tensor,为了加快速度，通过共享内存的方式
        # x_train_tensor = torch.FloatTensor(x).to(self.device)
        # y_train_tensor = torch.from_numpy(y).to(self.device)
        # # print(x_train_tensor.dtype,y_train_tensor.dtype)
        # train_ds = data.TensorDataset(x_train_tensor,y_train_tensor)
        # self.train_dl = data.DataLoader(train_ds,batch_size=self.batch_size,shuffle=True)
        
        # self.iteration  = self.epochs*(len(self.train_dl))
        # logger.info("FedPerClient process_data: iteration:{}\t".format(self.iteration))
        self.process_data_support_and_query()
        # return self.train_dl

    def process_data_support_and_query(self):
        """注意，这里将support&&query set都设置为一个batch，相当于one batch梯度下降。
        """        
        x,y = self.datapair
        assert len(x) == len(y)
        sq_point  = int(len(y)*self.support_ratio)
        # 将数据转换成tensor,为了加快速度，通过共享内存的方式
        x_support_tensor = torch.FloatTensor(x[:sq_point]).to(self.device)
        y_support_tensor = torch.from_numpy(y[:sq_point]).to(self.device)
        x_query_tensor = torch.FloatTensor(x[sq_point:]).to(self.device)
        y_query_tensor = torch.from_numpy(y[sq_point:]).to(self.device)
        # print(x_train_tensor.dtype,y_train_tensor.dtype)
        support_ds = data.TensorDataset(x_support_tensor,y_support_tensor)
        self.support_dl = data.DataLoader(support_ds,batch_size=self.batch_size,shuffle=True)
        query_ds  = data.TensorDataset(x_query_tensor,y_query_tensor)
        self.query_dl = data.DataLoader(query_ds,batch_size=len(query_ds),shuffle=True)
        self.iteration  = self.epochs*(len(self.support_dl))
        logger.info("FedPerClient process_data_support_and_query: iteration_support_dl:{}\t".format(self.iteration))
        return self.support_dl,self.query_dl

    # base_layer load model
    def load_model_state_dict_base_layer(self,model_state_dict):
        # 加载部分参数，还是传递部分参数。传递部分参数过来。加载的时候都是创建一个新的。
        new_par = copy.deepcopy(self.model.state_dict())
        for name in model_state_dict:
            new_par[name]= model_state_dict[name]
        self.model.load_state_dict(new_par)


    def get_model_state_dict_base_layer(self):
        state_dict =  self.model.state_dict()
        base_layer_state_dict = {key:state_dict[key] for i,key in enumerate(state_dict)  if i < self.base_layer}
        return base_layer_state_dict
   



   # 用来训练当前的本地模型。每次训练返回损失值running_loss和running_accuracy
    def train_per(self):
        self.model.train()
        loss_sum = 0.
        loss_times = 0.
        correct_sum = 0.
        total_size = 0        

        for epoch in range(self.epochs):
            # 使用support_set进行梯度下降
            for batch_x,batch_y in self.support_dl:
                # 正向传播
                outputs = self.model(batch_x)
                _,predicted = torch.max(outputs,1)
                correct_sum += (predicted == batch_y).sum().item()
                total_size += len(batch_y)
                # 计算loss
                loss = self.loss_func(outputs,batch_y.long())
                loss_sum +=loss.item()
                loss_times+=1
                # 反向传播
                self.optimizer.zero_grad()
                loss.backward()

                # 梯度下降
                self.optimizer.step()
        return loss_sum/loss_times,100*correct_sum/total_size
    

    # 对本地的个性化层进行自适应训练，并在测试集上运行查看结果
    def test_per(self):
        """生成一个专门用来测试的模型，而不影响原来的梯度下降过程。
        fedamp虽然完全相似。但是fedamp在测试的时候没有本地适应下降，不会破坏模型。
        所以这个方法需要生成一个新的模型。专门用来测试。

        Returns:
            [type]: [description]
        """        
        # 复制一个模型
        test_model = copy.deepcopy(self.model)

        test_model.train()
        # 冻结base_layer层，训练local层。这里可以进行优化，将这样的两个优化器直接放到init中。所有的参数require_grade都是true。但是optimizer只更新部分参数。
        for i,parm in enumerate(test_model.parameters()):
            if(i<self.base_layer):
                parm.requires_grad_(False)
            else:
                parm.requires_grad=True
        test_optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, test_model.parameters()),lr=0.01)

        for epoch in range(self.epochs):
            # 使用support_set进行梯度下降
            for batch_x,batch_y in self.support_dl:
                # 正向传播
                outputs = test_model(batch_x)

                # 计算loss
                loss = self.loss_func(outputs,batch_y.long())
                # 反向传播
                test_optimizer.zero_grad()
                loss.backward()

                # 梯度下降
                test_optimizer.step()
        
        # 解冻所有的层
        for parm in test_model.parameters():
            parm.requires_grad=True
        test_optimizer = torch.optim.SGD(self.model.parameters(),lr=0.01)

        test_model.eval()
        correct = 0
        total =0
        with torch.no_grad():
            for features,labels in self.query_dl:
                outputs = test_model(features)
                # print(outputs.shape)
                _,predicted = torch.max(outputs,1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        accuracy = 100*correct/total
        return accuracy

